package org.example.cases.sf10;

import com.antgroup.geaflow.api.function.base.FilterFunction;
import com.antgroup.geaflow.api.function.base.MapFunction;
import com.antgroup.geaflow.api.graph.PGraphWindow;
import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowSource;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowStream;
import com.antgroup.geaflow.api.window.impl.AllWindow;
import com.antgroup.geaflow.common.config.Configuration;
import com.antgroup.geaflow.common.tuple.Tuple;
import com.antgroup.geaflow.model.graph.edge.EdgeDirection;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.edge.impl.ValueEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import com.antgroup.geaflow.model.graph.vertex.impl.ValueLabelVertex;
import com.antgroup.geaflow.pipeline.Pipeline;
import com.antgroup.geaflow.pipeline.task.PipelineTask;
import com.antgroup.geaflow.view.GraphViewBuilder;
import com.antgroup.geaflow.view.IViewDesc;
import com.antgroup.geaflow.view.graph.GraphViewDesc;
import org.example.config.JobConfig;
import org.example.constants.Constants;
import org.example.function.AbstractVcFunc;
import org.example.function.CSVFileSinkV2;
import org.example.function.CSVFileSourceV2;
import org.example.model.LongStringResult;
import org.example.util.FileUtils;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;

public class Case1Pipeline {

    private final String snapshotPath;

    private final Pipeline pipeline;

    public Case1Pipeline(String snapshotPath, Pipeline pipeline) {
        this.snapshotPath = snapshotPath;
        this.pipeline = pipeline;
    }

    public void submit() {

        pipeline.submit((PipelineTask) pipelineTaskCxt -> {
            Configuration conf = pipelineTaskCxt.getConfig();

            PWindowSource<IVertex<Long, Double>> accountVertices = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, "Account.csv"),
                            Collections.singletonList("accountId"),
                            (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex =
                                        new ValueLabelVertex<>(
                                                Long.valueOf(line[0]),
                                                null,
                                                Constants.VERTEX_ACCOUNT
                                        );
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SOURCE_PARALLELISM));

            PWindowSource<IVertex<Long, Double>> loanVertices = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, "Loan.csv"),
                            Arrays.asList("loanId", "loanAmount"),
                            (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex = new ValueLabelVertex<>(
                                        -Long.parseLong(line[0]),
                                        Double.parseDouble(line[1]),
                                        Constants.VERTEX_LOAN
                                );
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SOURCE_PARALLELISM));

            PWindowSource<IVertex<Long, Double>> personVertices = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, "Person.csv"),
                            Collections.singletonList("personId"),
                            (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex = new ValueLabelVertex<>(
                                        Long.parseLong(line[0]),
                                        0.0,
                                        Constants.VERTEX_PERSON
                                );
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SOURCE_PARALLELISM));

            PWindowStream<IVertex<Long, Double>> unionVertexStream = accountVertices
                    .union(personVertices)
                    .withParallelism(conf.getInteger(JobConfig.UNION_PARALLELISM))
                    .union(loanVertices)
                    .withParallelism(conf.getInteger(JobConfig.UNION_PARALLELISM));

            PWindowSource<IEdge<Long, Double>> personOwnAccountEdges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, "PersonOwnAccount.csv"),
                            Arrays.asList("personId", "accountId"),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        Long.parseLong(line[1]),
                                        Long.parseLong(line[0]),
                                        null,
                                        EdgeDirection.IN
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));

            PWindowSource<IEdge<Long, Double>> loanDepositAccountEdges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, "LoanDepositAccount.csv"),
                            Arrays.asList("loanId", "accountId", "amount"),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        -Long.parseLong(line[0]),
                                        Long.parseLong(line[1]),
                                        Double.parseDouble(line[2])
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));

            PWindowSource<IEdge<Long, Double>> accountTransferAccountEdges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, "AccountTransferAccount.csv"),
                            Arrays.asList("fromId", "toId", "amount"),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        Long.parseLong(line[0]),
                                        Long.parseLong(line[1]),
                                        Double.parseDouble(line[2])
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));

            PWindowStream<IEdge<Long, Double>> unionEdgeStream = personOwnAccountEdges
                    .union(loanDepositAccountEdges)
                    .withParallelism(conf.getInteger(JobConfig.UNION_PARALLELISM))
                    .union(accountTransferAccountEdges)
                    .withParallelism(conf.getInteger(JobConfig.UNION_PARALLELISM));

            int iterationParallelism = conf.getInteger(JobConfig.ITERATOR_PARALLELISM);
            GraphViewDesc graphViewDesc = GraphViewBuilder
                    .createGraphView("case1")
                    .withShardNum(iterationParallelism)
                    .withBackend(IViewDesc.BackendType.Memory)
                    .build();

            PGraphWindow<Long, Double, Double> graphWindow =
                    pipelineTaskCxt.buildWindowStreamGraph(unionVertexStream, unionEdgeStream, graphViewDesc);

            graphWindow.compute(new Case1VCAlgorithm(4))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Double>>) record ->
                            Constants.VERTEX_PERSON.equals(((ValueLabelVertex<Long, Double>) record).getLabel()) && record.getValue() > 0)
                    .withParallelism(conf.getInteger(JobConfig.FILTER_PARALLELISM))
                    .map((MapFunction<IVertex<Long, Double>, LongStringResult>) value ->
                            new LongStringResult(value.getId(), String.format("%.2f", value.getValue() / 100_000_000)))
                    .withParallelism(conf.getInteger(JobConfig.MAP_PARALLELISM))
                    .sink(new CSVFileSinkV2<>("result1.csv"))
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));
        });
    }

    public static class Case1VCAlgorithm extends VertexCentricCompute<Long, Double, Double, Tuple<Long, Double>> {

        public Case1VCAlgorithm(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricCombineFunction<Tuple<Long, Double>> getCombineFunction() {
            return null;
        }

        @Override
        public VertexCentricComputeFunction<Long, Double, Double, Tuple<Long, Double>> getComputeFunction() {
            return new Case4VertexCentricComputeFunction();
        }
    }

    public static class Case4VertexCentricComputeFunction extends AbstractVcFunc<Long, Double, Double, Tuple<Long, Double>> {

        @Override
        public void compute(Long vertexId, Iterator<Tuple<Long, Double>> messageIterator) {
            ValueLabelVertex<Long, Double> vertex = (ValueLabelVertex<Long, Double>) this.context.vertex().get();
            switch ((int) (context.getIterationId() % 4)) {
                case 1:
                    if (Constants.VERTEX_LOAN.equals(vertex.getLabel())) {
                        List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                        if (!outEdges.isEmpty()) {
                            for (IEdge<Long, Double> outEdge : outEdges) {
                                Long targetId = outEdge.getTargetId();
                                Double loanAmount = vertex.getValue();
                                this.context.sendMessage(targetId, Tuple.of(vertexId, loanAmount));
                            }
                        }
                    }
                case 2:
                    if (Constants.VERTEX_ACCOUNT.equals(vertex.getLabel())) {
                        List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
                        while (messageIterator.hasNext()) {
                            Tuple<Long, Double> message = messageIterator.next();
                            Long loanId = message.f0;
                            Double loanAmount = message.f1;
                            for (IEdge<Long, Double> outEdge : outEdges) {
                                Long targetId = outEdge.getTargetId();
                                this.context.sendMessage(targetId, Tuple.of(loanId, loanAmount));
                            }
                        }
                    }
                case 3:
                    List<IEdge<Long, Double>> inEdges = this.context.edges().getInEdges();
                    while (messageIterator.hasNext()) {
                        Tuple<Long, Double> message = messageIterator.next();
                        Long loanId = message.f0;
                        Double loanAmount = message.f1;
                        for (IEdge<Long, Double> inEdge : inEdges) {
                            Long targetId = inEdge.getTargetId();
                            this.context.sendMessage(targetId, Tuple.of(loanId, loanAmount));
                        }
                    }
                case 0:
                    if (Constants.VERTEX_PERSON.equals(vertex.getLabel())) {
                        double personLoanSum = 0.0;
                        Set<Long> loans = new HashSet<>();
                        while (messageIterator.hasNext()) {
                            Tuple<Long, Double> message = messageIterator.next();
                            long loanId = message.f0;
                            if (!loans.contains(loanId)) {
                                double loanAmount = message.f1;
                                personLoanSum += loanAmount;
                                loans.add(loanId);
                            }
                        }
                        this.context.setNewVertexValue(personLoanSum);
                    }
            }
        }
    }
}
